// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/common/compression/compression.h>
#include <pollux/common/config/config.h>
#include <pollux/dwio/common/data_buffer.h>
#include <pollux/dwio/common/file_sink.h>
#include <pollux/dwio/common/flush_policy.h>
#include <pollux/dwio/common/options.h>
#include <pollux/dwio/common/writer.h>
#include <pollux/dwio/common/writer_factory.h>
#include <pollux/dwio/parquet/writer/arrow/types.h>
#include <pollux/dwio/parquet/writer/arrow/util/compression.h>
#include <pollux/vector/complex_vector.h>
#include <pollux/vector/arrow/bridge.h>

namespace kumo::pollux::parquet {
    using kumo::pollux::parquet::arrow::util::CodecOptions;

    class ArrowDataBufferSink;

    struct ArrowContext;

    class DefaultFlushPolicy : public dwio::common::FlushPolicy {
    public:
        DefaultFlushPolicy()
            : rowsInRowGroup_(1'024 * 1'024), bytesInRowGroup_(128 * 1'024 * 1'024) {
        }

        DefaultFlushPolicy(uint64_t rowsInRowGroup, int64_t bytesInRowGroup)
            : rowsInRowGroup_(rowsInRowGroup), bytesInRowGroup_(bytesInRowGroup) {
        }

        bool shouldFlush(
            const dwio::common::StripeProgress &stripeProgress) override {
            return stripeProgress.stripeRowCount >= rowsInRowGroup_ ||
                   stripeProgress.stripeSizeEstimate >= bytesInRowGroup_;
        }

        void onClose() override {
            // No-op
        }

        uint64_t rowsInRowGroup() const {
            return rowsInRowGroup_;
        }

        int64_t bytesInRowGroup() const {
            return bytesInRowGroup_;
        }

    private:
        const uint64_t rowsInRowGroup_;
        const int64_t bytesInRowGroup_;
    };

    class LambdaFlushPolicy : public DefaultFlushPolicy {
    public:
        explicit LambdaFlushPolicy(
            uint64_t rowsInRowGroup,
            int64_t bytesInRowGroup,
            std::function<bool()> lambda)
            : DefaultFlushPolicy(rowsInRowGroup, bytesInRowGroup) {
            lambda_ = std::move(lambda);
        }

        virtual ~LambdaFlushPolicy() override = default;

        bool shouldFlush(
            const dwio::common::StripeProgress &stripeProgress) override {
            return lambda_() || DefaultFlushPolicy::shouldFlush(stripeProgress);
        }

    private:
        std::function<bool()> lambda_;
    };

    struct WriterOptions : public dwio::common::WriterOptions {
        bool enableDictionary = true;
        int64_t dataPageSize = 1'024 * 1'024;
        int64_t dictionaryPageSizeLimit = 1'024 * 1'024;
        // Growth ratio passed to ArrowDataBufferSink. The default value is a
        // heuristic borrowed from
        // melon/FBVector(https://github.com/facebook/melon/blob/main/melon/docs/FBVector.md#memory-handling).
        double bufferGrowRatio = 1.5;

        arrow::Encoding::type encoding = arrow::Encoding::PLAIN;

        std::shared_ptr<CodecOptions> codecOptions;
        std::unordered_map<std::string, common::CompressionKind>
        columnCompressionsMap;

        /// Timestamp unit for Parquet write through Arrow bridge.
        /// Default if not specified: TimestampPrecision::kNanoseconds (9).
        std::optional<TimestampPrecision> parquetWriteTimestampUnit;
        /// Timestamp time zone for Parquet write through Arrow bridge.
        std::optional<std::string> parquetWriteTimestampTimeZone;
        bool writeInt96AsTimestamp = false;
        std::optional<bool> useParquetDataPageV2;

        // Parsing session and hive configs.

        // This isn't a typo; session and hive connector config names are different
        // ('_' vs '-').
        static constexpr const char *kParquetSessionWriteTimestampUnit =
                "hive.parquet.writer.timestamp_unit";
        static constexpr const char *kParquetHiveConnectorWriteTimestampUnit =
                "hive.parquet.writer.timestamp-unit";
        static constexpr const char *kParquetSessionDataPageVersion =
                "hive.parquet.writer.datapage_version";
        static constexpr const char *kParquetHiveConnectorDataPageVersion =
                "hive.parquet.writer.datapage-version";

        // Process hive connector and session configs.
        void processConfigs(
            const config::ConfigBase &connectorConfig,
            const config::ConfigBase &session) override;
    };

    // Writes Pollux vectors into  a DataSink using Arrow Parquet writer.
    class Writer : public dwio::common::Writer {
    public:
        // Constructs a writer with output to 'sink'. A new row group is
        // started every 'rowsInRowGroup' top level rows. 'pool' is used for
        // temporary memory. 'properties' specifies Parquet-specific
        // options. 'schema' specifies the file's overall schema, and it is always
        // non-null.
        Writer(
            std::unique_ptr<dwio::common::FileSink> sink,
            const WriterOptions &options,
            std::shared_ptr<memory::MemoryPool> pool,
            RowTypePtr schema);

        Writer(
            std::unique_ptr<dwio::common::FileSink> sink,
            const WriterOptions &options,
            RowTypePtr schema);

        ~Writer() override = default;

        static bool isCodecAvailable(common::CompressionKind compression);

        // Appends 'data' into the writer.
        void write(const VectorPtr &data) override;

        void flush() override;

        // Forces a row group boundary before the data added by next write().
        void newRowGroup(int32_t numRows);

        bool finish() override {
            return true;
        }

        // Closes 'this', After close, data can no longer be added and the completed
        // Parquet file is flushed into 'sink' provided at construction. 'sink' stays
        // live until destruction of 'this'.
        void close() override;

        void abort() override;

    private:
        // Sets the memory reclaimers for all the memory pools used by this writer.
        void setMemoryReclaimers();

        // Pool for 'stream_'.
        std::shared_ptr<memory::MemoryPool> pool_;
        std::shared_ptr<memory::MemoryPool> generalPool_;

        // Temporary Arrow stream for capturing the output.
        std::shared_ptr<ArrowDataBufferSink> stream_;

        std::shared_ptr<ArrowContext> arrowContext_;

        std::unique_ptr<DefaultFlushPolicy> flushPolicy_;

        const RowTypePtr schema_;

        ArrowOptions options_{.flattenDictionary = true, .flattenConstant = true};

        // Whether to write Int96 timestamps in Arrow Parquet write.
        bool writeInt96AsTimestamp_;
    };

    class ParquetWriterFactory : public dwio::common::WriterFactory {
    public:
        ParquetWriterFactory() : WriterFactory(dwio::common::FileFormat::PARQUET) {
        }

        std::unique_ptr<dwio::common::Writer> createWriter(
            std::unique_ptr<dwio::common::FileSink> sink,
            const std::shared_ptr<dwio::common::WriterOptions> &options) override;

        std::unique_ptr<dwio::common::WriterOptions> createWriterOptions() override;
    };
} // namespace kumo::pollux::parquet
